from typing import Dict, List, Tuple, Any
import networkx as nx
import numpy as np
from dataclasses import dataclass
from enum import Enum
from sklearn.cluster import AgglomerativeClustering, KMeans
from collections import defaultdict

class SpaceType(Enum):
    PERCEPTION = "perception"
    PLANNING = "planning"
    REASONING = "reasoning"
    ACTION = "action"

@dataclass
class WisdomNode:
    id: str
    space_type: SpaceType
    features: Dict[str, Any]
    level: int
    cluster_id: int = -1  # For hierarchical clustering
    importance_score: float = 0.0
    reward: float = 0.0
    probability: float = 1.0

class WisdomGraph:
    def __init__(self):
        self.perception_graph = nx.Graph()
        self.planning_graph = nx.Graph()
        self.reasoning_graph = nx.Graph()
        self.action_graph = nx.Graph()
        self.cross_space_edges = nx.DiGraph()
        
        # Cluster management
        self.perception_clusters = defaultdict(list)
        self.cluster_importance = {}
        
    def add_perception_node(self, node_id: str, features: Dict[str, Any], 
                           level: int, importance_score: float) -> None:
        """
        Add a node in perception space with support for hierarchical clustering
        """
        self.perception_graph.add_node(node_id, features=features, level=level, importance=importance_score)

    def add_planning_node(self, node_id: str, task_info: Dict[str, Any], level: int) -> None:
        """Add a node in planning space (task hierarchies)"""
        self.planning_graph.add_node(node_id, task_info=task_info, level=level)

    def add_reasoning_node(self, node_id: str, variables: Dict[str, Any], 
                          probability: float, level: int) -> None:
        """Add a node in reasoning space (Bayesian network nodes)"""
        self.reasoning_graph.add_node(node_id, variables=variables, probability=probability, level=level)

    def add_action_node(self, node_id: str, state_info: Dict[str, Any], 
                       reward: float, level: int) -> None:
        """Add a node in action space (Q-learning states)"""
        self.action_graph.add_node(node_id, state_info=state_info, reward=reward, level=level)

    def add_edge_within_space(self, space_type: SpaceType, source_id: str, 
                            target_id: str, weight: float = 1.0) -> None:
        """Add an edge within a specific space"""
        graph = getattr(self, f"{space_type.value}_graph")
        if source_id in graph and target_id in graph:
            graph.add_edge(source_id, target_id, weight=weight)

    def add_cross_space_edge(self, source_id: str, target_id: str, 
                            source_space: SpaceType, target_space: SpaceType) -> None:
        """Connect nodes across different spaces"""
        self.cross_space_edges.add_edge(
            source_id, 
            target_id, 
            source_space=source_space,
            target_space=target_space
        )

    def cluster_perception_nodes(self, n_clusters: int = 5) -> None:
        """
        Group related perceptual features into clusters using hierarchical clustering
        """
        if len(self.perception_graph.nodes) < n_clusters:
            return
        
        # Simple clustering based on importance scores
        importance_scores = [data.get('importance', 0) 
                           for _, data in self.perception_graph.nodes(data=True)]
        importance_scores = np.array(importance_scores).reshape(-1, 1)
        
        kmeans = KMeans(n_clusters=n_clusters)
        kmeans.fit(importance_scores)
        
        # Assign clusters to nodes
        for (node, data), cluster in zip(self.perception_graph.nodes(data=True), kmeans.labels_):
            self.perception_graph.nodes[node]['cluster'] = cluster

    def calculate_cluster_importance(self) -> None:
        """
        Calculate importance scores for each cluster based on node importance
        """
        if not any('cluster' in data for _, data in self.perception_graph.nodes(data=True)):
            return
            
        cluster_scores = {}
        for node, data in self.perception_graph.nodes(data=True):
            cluster = data.get('cluster')
            if cluster is not None:
                if cluster not in cluster_scores:
                    cluster_scores[cluster] = []
                cluster_scores[cluster].append(data.get('importance', 0))
        
        self.cluster_importance = {
            cluster: sum(scores)/len(scores) 
            for cluster, scores in cluster_scores.items()
        }

    def prune_perception_space(self, cluster_threshold: float, 
                             node_threshold: float) -> None:
        """
        Implement hierarchical pruning for perception space
        """
        # First calculate cluster importance
        self.calculate_cluster_importance()
        
        nodes_to_remove = []
        
        # Prune based on cluster importance
        for cluster_id, importance in self.cluster_importance.items():
            if importance < cluster_threshold:
                # Remove all nodes in low-importance clusters
                nodes_to_remove.extend(self.perception_clusters[cluster_id])
            else:
                # For important clusters, still prune individual low-importance nodes
                for node_id in self.perception_clusters[cluster_id]:
                    node_importance = self.perception_graph.nodes[node_id]['data'].importance_score
                    if node_importance < node_threshold:
                        nodes_to_remove.append(node_id)

        # Remove nodes and update clusters
        self.perception_graph.remove_nodes_from(nodes_to_remove)
        
        # Update cluster assignments
        self.perception_clusters = defaultdict(list)
        for node_id, data in self.perception_graph.nodes(data=True):
            cluster_id = data['data'].cluster_id
            if cluster_id >= 0:
                self.perception_clusters[cluster_id].append(node_id)

    def update_perception_importance(self, node_id: str, 
                                   importance_delta: float) -> None:
        """
        Update importance score for a perception node based on its contribution
        """
        if node_id in self.perception_graph:
            node = self.perception_graph.nodes[node_id]['data']
            node.importance_score = max(0.0, min(1.0, 
                node.importance_score + importance_delta))

    def get_cluster_nodes(self, cluster_id: int) -> List[str]:
        """
        Get all nodes belonging to a specific cluster
        """
        return self.perception_clusters.get(cluster_id, [])

    def get_cluster_importance(self, cluster_id: int) -> float:
        """
        Get the importance score for a specific cluster
        """
        return self.cluster_importance.get(cluster_id, 0.0)

    def prune_planning_space(self, efficiency_threshold: float) -> None:
        """Prune inefficient subtasks in planning space"""
        nodes_to_remove = []
        for node_id, data in self.planning_graph.nodes(data=True):
            if data['data'].features.get('efficiency', 0) < efficiency_threshold:
                nodes_to_remove.append(node_id)
        self.planning_graph.remove_nodes_from(nodes_to_remove)

    def prune_reasoning_space(self, probability_threshold: float) -> None:
        """Prune less influential nodes in reasoning space"""
        nodes_to_remove = []
        for node_id, data in self.reasoning_graph.nodes(data=True):
            if data['data'].probability < probability_threshold:
                nodes_to_remove.append(node_id)
        self.reasoning_graph.remove_nodes_from(nodes_to_remove)

    def prune_action_space(self, reward_threshold: float) -> None:
        """Prune low-reward actions"""
        nodes_to_remove = []
        for node_id, data in self.action_graph.nodes(data=True):
            if data['data'].reward < reward_threshold:
                nodes_to_remove.append(node_id)
        self.action_graph.remove_nodes_from(nodes_to_remove)

    def update_q_value(self, state_id: str, action_id: str, reward: float, 
                      learning_rate: float = 0.1) -> None:
        """Update Q-values in action space"""
        if state_id in self.action_graph and action_id in self.action_graph:
            current_q = self.action_graph.nodes[state_id]['data'].reward
            new_q = current_q + learning_rate * (reward - current_q)
            self.action_graph.nodes[state_id]['data'].reward = new_q

    def get_highest_reward_action(self, state_id: str) -> str:
        """Get the action with highest reward for a given state"""
        if state_id not in self.action_graph:
            return None
        
        successors = list(self.action_graph.successors(state_id))
        if not successors:
            return None
            
        return max(
            successors,
            key=lambda x: self.action_graph.nodes[x]['data'].reward
        )
